import torch
from mindspore import Tensor, Parameter
from collections import defaultdict


# import argparse
#
# parser = argparse.ArgumentParser('Convert torch model to mindspore')
# parser.add_argument('-s', '--source', help='The path torch model', required=True)
# parser.add_argument('-t', '--target', default='mindspore.ckpt', help='The save path')
#

def to_mindspore(x: torch.Tensor, name: str):
    return Parameter(Tensor(x.cpu().detach().numpy()), name=name)


def convert(pth: dict, check=None):
    groups = defaultdict(list)
    for name, value in pth.items():
        s = name.split('.')
        prefix = '.'.join(s[:-1])
        groups[prefix].append((s[-1], value))
    save = {}
    for prefix, v in groups.items():
        if len(v) == 5:  # BatchNorm2d
            names = [x[0] for x in v]
            assert 'running_var' in names and 'running_mean' in names
            for x in v:
                if x[0] == 'weight':
                    name = (prefix + '.gamma').strip('.')
                    save[name] = to_mindspore(x[1], name)
                elif x[0] == 'bias':
                    name = (prefix + '.beta').strip('.')
                    save[name] = to_mindspore(x[1], name)
                elif x[0] == 'running_mean':
                    name = (prefix + '.moving_mean').strip('.')
                    save[name] = to_mindspore(x[1], name)
                elif x[0] == 'running_var':
                    name = (prefix + '.moving_variance').strip('.')
                    save[name] = to_mindspore(x[1], name)
        else:  # nn.Conv2d
            for x in v:
                name = '.'.join([prefix, x[0]])
                save[name] = to_mindspore(x[1], name)
    if check is not None:
        assert len(save) == len(check)
        for key, value in check.items():
            assert key in save
            cmp = save[key]
            assert cmp.shape == value.shape
            assert cmp.name == value.name
            assert cmp.dtype == value.dtype
    return save
